from keras.preprocessing.image import img_to_array
from keras.models import load_model
import numpy as np
import argparse
import imutils
import cv2
import tensorflow as tf 

def detect():
    
    detector = cv2.CascadeClassifier(cascade_path)
    
    if video_path == None:
        camera = cv2.VideoCapture(0)
    else:
        camera = cv2.VideoCapture(video_path)

    with tf.Session() as sess:

        saver = tf.train.import_meta_graph(model_path)
        ckpt = tf.train.get_checkpoint_state(checkpoint_path)
        saver.restore(sess,ckpt.all_model_checkpoint_paths[0])
        graph = tf.get_default_graph()
        X = graph.get_operation_by_name('X').outputs[0]
        prediction = tf.get_collection('prediction')

        while True:
            (grabbed,frame) = camera.read()

            if video_path and not grabbed:
                break
        
            frame_data = imutils.resize(frame,width = 1000)
            gray = cv2.cvtColor(frame_data,cv2.COLOR_BGR2GRAY)
            frameClone = frame_data.copy()
            rects = detector.detectMultiScale(gray,scaleFactor = 1.1,minNeighbors = 5,minSize = (30,30),flags = cv2.CASCADE_SCALE_IMAGE)
                
          
            for (fX,fY,fW,fH) in rects:
            
                roi = gray[fY:fY + fH,fX:fX + fW]
                roi = cv2.resize(roi,(28,28))
                roi = roi.astype("float") / 255.0
                roi = img_to_array(roi)
                roi = np.expand_dims(roi,axis = 0)

                smiling_score = sess.run(prediction[0][0][1],feed_dict = {X:roi})
                notSmiling_score = sess.run(prediction[0][0][0],feed_dict = {X:roi})

                label = "Smiling" if smiling_score > notSmiling_score else "Not Smiling"

                cv2.putText(frameClone,label,(fX,fY - 10),cv2.FONT_HERSHEY_SIMPLEX,0.45,(0,0,255),2)
                cv2.rectangle(frameClone,(fX,fY),(fX + fW,fY + fH),(0,0,255),2)
            cv2.imshow("Face",frameClone)

            if cv2.waitKey(1) & 0xFF == ord("q"):
                break

    camera.release()
    cv2.destroyAllWindows()
    


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("--cascade",required = True,help = "path to where the face cascade resides")
    parser.add_argument("--model",required = True,help = "path to pre-trained smile detector CNN")
    parser.add_argument("--checkpoint",required = True,help = "path to checkpoint for pre-trained smile detector CNN")
    parser.add_argument("--video",help = "path to  the video file")

    args = parser.parse_args()

    cascade_path = args.cascade
    model_path = args.model
    video_path = args.video
    checkpoint_path = args.checkpoint

    detect()
    
